/*
 * Copyright 2021 TiKV Project Authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package com.webank.tikv;

import com.google.common.base.Strings;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import com.google.protobuf.ByteString;
import com.webank.tikv.common.TiTimestamp;
import com.webank.tikv.exception.GrpcException;
import com.webank.tikv.exception.InternalException;
import com.webank.tikv.kvproto.Metapb;
import com.webank.tikv.kvproto.PDGrpc;
import com.webank.tikv.kvproto.Pdpb;
import com.webank.tikv.operation.PDErrorHandler;
import com.webank.tikv.util.BackOffer;
import com.webank.tikv.util.ChannelFactory;
import com.webank.tikv.util.ConcreteBackOffer;
import com.webank.tikv.util.UriUtils;
import io.grpc.ManagedChannel;
import io.grpc.health.v1.HealthCheckRequest;
import io.grpc.health.v1.HealthCheckResponse;
import io.grpc.health.v1.HealthGrpc;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import java.net.URI;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Supplier;
import java.util.stream.Collectors;

import static com.webank.tikv.common.PDError.buildFromPdpbError;
import static com.webank.tikv.operation.PDErrorHandler.PD_RESPONSE_ERROR_FUNCTION;
import static com.webank.tikv.util.BackOffer.BoType.BoCheckHealth;

public class PDClient extends AbstractGRPCClient<PDGrpc.PDBlockingStub, PDGrpc.PDFutureStub> {

    private static final Logger LOGGER = LogManager.getLogger(PDClient.class);

    private static final long MIN_TRY_UPDATE_DURATION = 50;

    private final List<URI> pdAddrs = new CopyOnWriteArrayList<>();

    private Pdpb.RequestHeader header;
    private Pdpb.TsoRequest tsoReq;

    private PDGrpc.PDBlockingStub blockingStub;
    private PDGrpc.PDFutureStub asyncStub;
    private String leaderAddr;

    private long lastUpdateLeaderTime;
    private ScheduledExecutorService pbExecutor;
    private final AtomicBoolean updateLeaderNotify = new AtomicBoolean();

    public PDClient(TiConfiguration conf, ChannelFactory channelFactory) {
        super(conf, channelFactory);
        this.pdAddrs.addAll(conf.getPdAddrs());

        initCluster();
        this.blockingStub = getBlockingStub();
        this.asyncStub = getAsyncStub();
    }

    private void initCluster() {
        // shuffle PD addresses so that clients call getMembers from different PD
        Collections.shuffle(pdAddrs);

        // The first request may cost too much latency
        long originTimeout = conf.getTimeout();
        timeout = conf.getPdFirstGetMemberTimeout();
        Pdpb.GetMembersResponse resp = null;
        for (URI u : pdAddrs) {
            resp = getMembers(defaultBackOffer(), u);
            if (resp != null) {
                break;
            }
        }
        if (resp == null) {
            LOGGER.warn("cloud not get pd leader member with: {}", pdAddrs);
            throw new InternalException("cloud not get pd leader member");
        }

        timeout = originTimeout;
        long clusterId = resp.getHeader().getClusterId();
        header = Pdpb.RequestHeader.newBuilder().setClusterId(clusterId).build();
        tsoReq = Pdpb.TsoRequest.newBuilder().setHeader(header).setCount(1).build();

        List<URI> uris = resp.getMembersList()
            .stream()
            .map(mem -> UriUtils.addrToUri(mem.getClientUrls(0)))
            .collect(Collectors.toList());
        pdAddrs.clear();
        pdAddrs.addAll(uris);

        String leaderUrl = resp.getLeader().getClientUrls(0);
        leaderUrl = UriUtils.format(UriUtils.addrToUri(leaderUrl));
        createLeaderClient(leaderUrl);

        pbExecutor = Executors.newSingleThreadScheduledExecutor(
            new ThreadFactoryBuilder().setNameFormat("pd-bg-%d").setDaemon(true).build());
        long period = conf.getHealthPdCheckPeriod();
        pbExecutor.scheduleAtFixedRate(() -> {
            try {
                tryUpdateLeader();
            } catch (Exception e) {
                LOGGER.warn("update pd leader failed, ex:{}", e.getMessage());
            }
        }, period, period, TimeUnit.MILLISECONDS);
        LOGGER.info("init tikv cluster with pd: {}", pdAddrs);
    }

    private void tryUpdateLeader() {
        for (URI url : pdAddrs) {
            BackOffer backOffer = this.probeBackOffer();
            // since resp is null, we need update leader's address by walking through all pd server.
            Pdpb.GetMembersResponse resp = getMembers(backOffer, url);
            if (resp == null) {
                continue;
            }
            List<URI> uris = resp.getMembersList().stream()
                .map(mem -> UriUtils.addrToUri(mem.getClientUrls(0)))
                .collect(Collectors.toList());
            String leaderUrl = resp.getLeader().getClientUrlsList().get(0);
            leaderUrl = UriUtils.format(UriUtils.addrToUri(leaderUrl));

            // If leader is not change but becomes available, we can cancel follower forward.
            if (checkPdHealth(backOffer, leaderUrl) && trySwitchLeader(leaderUrl)) {
                LOGGER.info("success update pd leader:{} from: {}", leaderUrl, url);
                if (!uris.equals(pdAddrs)) {
                    pdAddrs.clear();
                    pdAddrs.addAll(uris);
                }
                lastUpdateLeaderTime = System.currentTimeMillis();
                return;
            }
        }
        lastUpdateLeaderTime = System.currentTimeMillis();
        LOGGER.warn("fail to update pd leader from: {}", pdAddrs);
    }

    // pd返回错误报文后, 在error handler中调用
    // 异步提交刷新pd leader任务
    public void scheduleUpdateLeader() {
        if (updateLeaderNotify.compareAndSet(false, true)) {
            try {
                pbExecutor.submit(() -> {
                    try {
                        updateLeader();
                    } catch (Exception e) {
                        LOGGER.warn("update leader exception:{}", e.getMessage());
                    }
                    updateLeaderNotify.set(false);
                });
            } catch (RuntimeException e) {
                LOGGER.warn("submit update leader task fail:{}", e.getMessage());
                updateLeaderNotify.set(false);
            }
        }
    }

    public TiTimestamp getTimestamp(BackOffer backOffer) {
        Supplier<Pdpb.TsoRequest> request = () -> tsoReq;
        PDErrorHandler<Pdpb.TsoResponse> handler = new PDErrorHandler<>(r -> r.getHeader().hasError() ?
            buildFromPdpbError(r.getHeader().getError()) : null, this);
        Pdpb.TsoResponse resp = callWithRetry(backOffer, PDGrpc.getTsoMethod(), request, handler);
        Pdpb.Timestamp timestamp = resp.getTimestamp();
        return new TiTimestamp(timestamp.getPhysical(), timestamp.getLogical());
    }

    public Pair<Metapb.Region, Metapb.Peer> getRegionByKey(BackOffer backOffer, ByteString key) {
        Supplier<Pdpb.GetRegionRequest> request = () -> Pdpb.GetRegionRequest.newBuilder()
            .setHeader(header).setRegionKey(key).build();
        PDErrorHandler<Pdpb.GetRegionResponse> handler =
            new PDErrorHandler<>(PD_RESPONSE_ERROR_FUNCTION, this);
        Pdpb.GetRegionResponse resp = callWithRetry(backOffer, PDGrpc.getGetRegionMethod(), request, handler);
        return Pair.of(resp.getRegion(), resp.getLeader());
    }

    public Metapb.Peer getLeaderByRegionId(BackOffer backOffer, long regionId) {
        Supplier<Pdpb.GetRegionByIDRequest> request = () -> Pdpb.GetRegionByIDRequest.newBuilder()
            .setHeader(header).setRegionId(regionId).build();
        PDErrorHandler<Pdpb.GetRegionResponse> handler =
            new PDErrorHandler<>(PD_RESPONSE_ERROR_FUNCTION, this);
        Pdpb.GetRegionResponse resp = callWithRetry(backOffer, PDGrpc.getGetRegionByIDMethod(), request, handler);
        return resp.getLeader();
    }

    public List<Pdpb.Region> scanRegions(BackOffer backOffer, ByteString startKey, ByteString endKey, int limit) {
        // no need to backoff because ScanRegions is just for optimization
        // introduce a warm-up timeout for ScanRegions requests
        PDGrpc.PDBlockingStub stub = getBlockingStub()
            .withDeadlineAfter(conf.getWarmUpTimeout(), TimeUnit.MILLISECONDS);
        Pdpb.ScanRegionsRequest request = Pdpb.ScanRegionsRequest.newBuilder()
            .setHeader(header)
            .setStartKey(startKey)
            .setEndKey(endKey)
            .setLimit(limit)
            .build();
        Pdpb.ScanRegionsResponse resp = stub.scanRegions(request);
        if (resp == null) {
            return null;
        }
        return resp.getRegionsList();
    }

    public Metapb.Store getStore(BackOffer backOffer, long storeId) {
        Supplier<Pdpb.GetStoreRequest> request = () ->
            Pdpb.GetStoreRequest.newBuilder().setHeader(header).setStoreId(storeId).build();
        PDErrorHandler<Pdpb.GetStoreResponse> handler = new PDErrorHandler<>(r -> r.getHeader().hasError() ?
            buildFromPdpbError(r.getHeader().getError()) : null, this);
        Pdpb.GetStoreResponse resp = callWithRetry(backOffer, PDGrpc.getGetStoreMethod(), request, handler);
        return resp != null ? resp.getStore() : null;
    }

    public List<Metapb.Store> getAllStores(BackOffer backOffer) {
        Supplier<Pdpb.GetAllStoresRequest> request = () ->
            Pdpb.GetAllStoresRequest.newBuilder().setHeader(header).build();
        PDErrorHandler<Pdpb.GetAllStoresResponse> handler = new PDErrorHandler<>(r -> r.getHeader().hasError() ?
            buildFromPdpbError(r.getHeader().getError()) : null, this);
        Pdpb.GetAllStoresResponse response = callWithRetry(backOffer, PDGrpc.getGetAllStoresMethod(), request, handler);
        return response.getStoresList();
    }

    private Pdpb.GetMembersResponse getMembers(BackOffer backOffer, URI uri) {
        while (true) {
            backOffer.checkTimeout();
            try {
                ManagedChannel probChan = channelFactory.getChannel(UriUtils.format(uri));
                PDGrpc.PDBlockingStub stub =
                    PDGrpc.newBlockingStub(probChan).withDeadlineAfter(getTimeout(), TimeUnit.MILLISECONDS);
                Pdpb.GetMembersRequest request =
                    Pdpb.GetMembersRequest.newBuilder().setHeader(Pdpb.RequestHeader.getDefaultInstance()).build();
                Pdpb.GetMembersResponse resp = stub.getMembers(request);
                // check if the response contains a valid leader
                if (resp != null && resp.getLeader().getMemberId() == 0) {
                    return null;
                }
                return resp;
            } catch (Exception e) {
                LOGGER.warn("failed to get member from pd server:{}, ex:{}", uri, e.getMessage());
                backOffer.doBackOff(BackOffer.BoType.BoPDRPC, e);
            }
        }
    }

    // return whether the leader has changed to target address leaderUrl
    private synchronized boolean trySwitchLeader(String leaderUrl) {
        if (!Strings.isNullOrEmpty(leaderAddr) && leaderUrl.equals(leaderAddr)) {
            return true;
        }
        return createLeaderClient(leaderUrl);
    }

    private synchronized boolean createLeaderClient(String leaderUrl) {
        try {
            ManagedChannel clientChannel = channelFactory.getChannel(leaderUrl);
            this.blockingStub = PDGrpc.newBlockingStub(clientChannel);
            this.asyncStub = PDGrpc.newFutureStub(clientChannel);
            this.leaderAddr = leaderUrl;
            this.timeout = conf.getTimeout();
            LOGGER.info("Switched to new pd leader: {}", leaderUrl);
            return true;
        } catch (RuntimeException e) {
            return false;
        }
    }

    private synchronized void updateLeader() {
        if (System.currentTimeMillis() - lastUpdateLeaderTime < MIN_TRY_UPDATE_DURATION) {
            LOGGER.warn("skip update pd leader task");
            return;
        }
        for (URI url : pdAddrs) {
            BackOffer backOffer = this.probeBackOffer();
            // since resp is null, we need update leader's address by walking through all pd server.
            Pdpb.GetMembersResponse resp = getMembers(backOffer, url);
            if (resp == null || resp.getLeader().getClientUrlsList().isEmpty()) {
                continue;
            }

            String leaderUrl = resp.getLeader().getClientUrlsList().get(0);
            leaderUrl = UriUtils.format(UriUtils.addrToUri(leaderUrl));
            if (checkPdHealth(backOffer, leaderUrl) && createLeaderClient(leaderUrl)) {
                lastUpdateLeaderTime = System.currentTimeMillis();
                LOGGER.info("schedule update pd leader success");
                return;
            }
        }
        lastUpdateLeaderTime = System.currentTimeMillis();
        LOGGER.warn("schedule update pd leader fail");
    }

    @Override
    protected PDGrpc.PDBlockingStub getBlockingStub() {
        if (blockingStub == null) {
            throw new GrpcException("PDClient may not be initialized");
        }
        return blockingStub.withDeadlineAfter(getTimeout(), TimeUnit.MILLISECONDS);
    }

    @Override
    protected PDGrpc.PDFutureStub getAsyncStub() {
        if (asyncStub == null) {
            throw new GrpcException("PDClient may not be initialized");
        }
        return asyncStub.withDeadlineAfter(getTimeout(), TimeUnit.MILLISECONDS);
    }

    public Long getClusterId() {
        return header.getClusterId();
    }

    public List<URI> getPdAddrs() {
        return pdAddrs;
    }

    private BackOffer defaultBackOffer() {
        return ConcreteBackOffer.newCustomBackOff(BackOffer.PD_INFO_BACKOFF);
    }

    private BackOffer probeBackOffer() {
        int maxSleep = (int) getTimeout() * 2;
        return ConcreteBackOffer.newCustomBackOff(maxSleep);
    }

    private boolean checkPdHealth(BackOffer backOffer, String address) {
        while (true) {
            backOffer.checkTimeout();
            try {
                ManagedChannel channel = channelFactory.getChannel(address);
                HealthGrpc.HealthBlockingStub stub = HealthGrpc.newBlockingStub(channel)
                    .withDeadlineAfter(getTimeout(), TimeUnit.MILLISECONDS);
                HealthCheckRequest req = HealthCheckRequest.newBuilder().build();
                HealthCheckResponse resp = stub.check(req);
                return resp.getStatus() == HealthCheckResponse.ServingStatus.SERVING;
            } catch (Exception e) {
                LOGGER.warn("check health failed, addr:{}, ex:{}", address, e.getMessage());
                backOffer.doBackOff(BoCheckHealth, e);
            }
        }
    }

    public void close() {
        if (pbExecutor != null) {
            pbExecutor.shutdownNow();
        }
    }

}
